import os, cv2
import sys
sys.path.append('..')
import glob
import pandas as pd
import numpy as np
from pathlib import Path
from utils.helper import read_xml, find_nodes, change_node_text, indent, write_xml
from utils.json import json_to_image


def json2xml(img_path, dst_path):
    # print(json_path)
    label_path = str(Path(img_path).with_suffix('.json'))
    if not os.path.exists(label_path):
        label_path = str(Path(img_path).with_suffix('.png'))
        if not os.path.exists(label_path):
            return 
    
    if label_path.endswith('.json'):
        try:
            img_path = str(Path(label_path).with_suffix('.jpg'))
            if not  os.path.exists(img_path): return
            img = cv2.imread(img_path, 1)
            mask = json_to_image(img.shape, label_path).astype('uint8')
            shape = img.shape[:2]
        except:
            return 
    elif label_path.endswith('.png'):
        mask = cv2.imread(label_path, 0).astype('uint8')
        shape = mask.shape

    base_tree = read_xml("../datasets/base_example.xml")
    root = base_tree.getroot()
    anno_tree = read_xml("../datasets/anno_example.xml")
    folder_node = find_nodes(base_tree, "folder")
    filename_node = find_nodes(base_tree, "filename")
    path_node = find_nodes(base_tree, "path")
    width_node = find_nodes(base_tree, "size/width")
    height_node = find_nodes(base_tree, "size/height")
    depth_node = find_nodes(base_tree, "size/depth")
    change_node_text(folder_node, Path(label_path).parent.name)
    change_node_text(filename_node, Path(label_path).name)
    change_node_text(path_node, str(Path(label_path).parent))
    change_node_text(width_node, str(shape[1]))
    change_node_text(height_node, str(shape[0]))
    change_node_text(depth_node, str(3))

    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    x_list = []
    y_list = []
    radius_list = []
    for cnt in contours:
        (x,y), radius = cv2.minEnclosingCircle(cnt)
        x_list.append(x)
        y_list.append(y)
        radius_list.append(radius)
    if len(radius_list) > 0:
        # cv2.imwrite(os.path.join(pred_mask, names[0] + '.png'), mask)  # 预测出缺陷，则保存mask
        index = np.argmax(radius_list)
        # defect_size_diameter = 2 * np.nanmax(radius_list)
        x, y, w, h = cv2.boundingRect(contours[index])

        x_min, x_max, y_min, y_max = x, x+w, y, y+h
        # img_boxed = img_ori[y_min:y_max, x_min:x_max, ...]

        xmin_node = find_nodes(anno_tree, "bndbox/xmin")
        ymin_node = find_nodes(anno_tree, "bndbox/ymin")
        xmax_node = find_nodes(anno_tree, "bndbox/xmax")
        ymax_node = find_nodes(anno_tree, "bndbox/ymax")
        change_node_text(xmin_node, str(x_min))
        change_node_text(ymin_node, str(y_min))
        change_node_text(xmax_node, str(x_max))
        change_node_text(ymax_node, str(y_max))
        root.append(anno_tree.getroot())
        indent(root)

        code = Path(label_path).parent.name
        xml_name = Path(label_path).with_suffix('.xml').name
        dst = os.path.join(dst_path, code, xml_name)
        os.makedirs(os.path.join(dst_path, code), exist_ok=True)
        write_xml(base_tree, dst)

def main():

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    # root_p = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_data'
    # ps = []
    # for product in os.listdir(root_p):
    #     # for folder in os.listdir(os.path.join(root_p, product)):
    #     ps.append(os.path.join(root_p, product))
    #     # if not os.path.isdir(os.path.join(root_p, product)):
    #     #     continue
    #     # for folder in os.listdir(os.path.join(root_p, product)):
    #     #     ps.append(os.path.join(root_p, product, folder))
    # print(ps)
    
    # dst_root = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_xml'
    # ds = []
    # for p in ps:
    #     product = str(Path(p).name)
    #     # folder = str(Path(p).name)
    #     ds.append(os.path.join(dst_root, product))

    p_root = r'/data0/autorepair/AA_T7_T10/t10/optical_data'
    ps = []
    for product in os.listdir(p_root):
          ps.append(os.path.join(p_root, product))

    ds = None

    if ds is None:
        ds = [x+"_xml" for x in ps]

    for path, dst_path in zip(ps, ds):

        tmp = []
        for code in os.listdir(path):
            imgs = glob.glob(os.path.join(path, code, "**/*.jpg"), recursive=True)
            tmp.extend(imgs)

        df = pd.DataFrame()
        df['img_path'] = tmp

        if is_pandarallel:
            df.parallel_apply(lambda x:json2xml(img_path=x['img_path'], dst_path=dst_path), axis=1)
        else:
            df.apply(lambda x:json2xml(img_path=x['img_path'], dst_path=dst_path), axis=1)
        print()


    # #####################################################################################
    # try:
    #     from pandarallel import pandarallel
    #     pandarallel.initialize(progress_bar=True) 
    #     print('Use multi threading !')
    #     is_pandarallel = True
    # except:
    #     print('Use single threading !')
    #     is_pandarallel = False

    # path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_data'
    # dst_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t10/optical_xml'

    # for product in os.listdir(path):
    #     print('product: ', product)
    #     for folder in os.listdir(os.path.join(path, product)):
    #         print('folder: ', folder)
    #         jsons = glob.glob(os.path.join(path, product, folder, "*/*.json"))
    #         if len(jsons) == 0: continue
    #         df = pd.DataFrame()
    #         df['json_path'] = jsons

    #         dst_path_cur = os.path.join(dst_path, product, folder)

    #         if is_pandarallel:
    #             df.parallel_apply(lambda x:json2xml(json_path=x['json_path'], dst_path=dst_path_cur), axis=1)
    #         else:
    #             df.apply(lambda x:json2xml(json_path=x['json_path'], dst_path=dst_path_cur), axis=1)
    #         print()


if __name__=='__main__':
    main()
